#include <algorithm>
#include <string>

#include "sam.h"
#include "comb_types.hpp"
#include "util.hpp"

const unsigned int SNP_CAND_FILE_TYPE     = 1;
const unsigned int CALL_RATES_FILE_TYPE   = 2;
const unsigned int SUBS_TOTL_FILE_TYPE    = 3;
const unsigned int SNP_CALL_COL_FILE_TYPE = 4;
const unsigned int SNP_CALL_SEQ_FILE_TYPE = 5;
const unsigned int EXPRESSION_FILE_TYPE   = 6;

const unsigned int SNP_CAND_FILE_VERSION       = 2;
const unsigned int CALL_RATES_COL_FILE_VERSION = 1;
const unsigned int CALL_RATES_SEQ_FILE_VERSION = 2;
const unsigned int SUBS_TOTL_FILE_VERSION      = 2;
const unsigned int SNP_CALL_COL_FILE_VERSION   = 4;
const unsigned int SNP_CALL_SEQ_FILE_VERSION   = 4;
const unsigned int EXPRESSION_FILE_VERSION     = 1;

std::map<char, std::map<char, char> > next_color;
std::map<char,int> iswap;
std::map<char, int> cswap;
std::map<char, std::map<char, int> > cswap2;
std::vector<SnpModel> haploid_models;
// a model for every base. A, C, G, T
std::vector<std::vector<SnpModel> > diploid_models;

void error(const std::string &msg)
{
    std::cerr << "ERROR: " << msg << std::endl;
    exit(1);
}

std::map<std::string, unsigned int> build_fasta_index(const std::string &fname)
{
    std::map<std::string, unsigned int> idx;
    std::fstream f(fname.c_str());
    std::string t, seqname="";
    unsigned int seqlen = 0;
    
    if(!f)
        error("failed to open the fasta file: " + fname);

    while(std::getline(f, t))
    {
        if(t[0] == '>')
        {
            if(seqlen != 0)
            {
                if(idx.find(seqname) != idx.end())
                    error("while parsing the fasta file: " + fname + "duplicate sequence names where encountered. |" + seqname + "| appears twice.");
                    
                idx[seqname] = seqlen;
            }
            
            std::stringstream ss(t.substr(1));
            ss >> seqname;
            seqlen = 0;
        }
        else
            seqlen += t.size();
    }

    if(seqlen != 0)
    {
        if(idx.find(seqname) != idx.end())
            error("while parsing the fasta file: " + fname + "duplicate sequence names where encountered. |" + seqname + "| appears twice.");
        idx[seqname] = seqlen;
    }

    return idx;
}

// reverse complement of a DNA base.
char rc(char c)
{
    if(c=='A'){ return 'T'; };
    if(c=='C'){ return 'G'; };
    if(c=='G'){ return 'C'; };
    if(c=='T'){ return 'A'; };
    return c;
}

// return a new string which is the reverse complement of the given DNA sequence.
std::string rc(std::string seq)
{
    std::string rev(seq);
    std::reverse(rev.begin(), rev.end());
    for(unsigned int i=0; i<rev.size(); i++)
        rev[i] = rc(rev[i]);
    return rev;
}


// We need to call this for every program that uses any global variables.  This is where they are all filled.
void init_globals()
{
    next_color['A']['0'] = 'A';
    next_color['A']['0'] = 'A';
    next_color['A']['1'] = 'C';
    next_color['A']['2'] = 'G';
    next_color['A']['3'] = 'T';
    next_color['C']['0'] = 'C';
    next_color['C']['1'] = 'A';
    next_color['C']['2'] = 'T';
    next_color['C']['3'] = 'G';
    next_color['G']['0'] = 'G';
    next_color['G']['1'] = 'T';
    next_color['G']['2'] = 'A';
    next_color['G']['3'] = 'C';
    next_color['T']['0'] = 'T';
    next_color['T']['1'] = 'G';
    next_color['T']['2'] = 'C';
    next_color['T']['3'] = 'A';

    iswap['0'] = 0;
    iswap['1'] = 1;
    iswap['2'] = 2;
    iswap['3'] = 3;

    cswap['A'] = 0;
    cswap['C'] = 1;
    cswap['G'] = 2;
    cswap['T'] = 3;
    
    cswap[0] = 'A';
    cswap[1] = 'C';
    cswap[2] = 'G';
    cswap[3] = 'T';

    cswap2['A']['A'] = 0;
    cswap2['A']['M'] = 1;
    cswap2['A']['R'] = 2;
    cswap2['A']['W'] = 3;
    cswap2['C']['C'] = 0;
    cswap2['C']['M'] = 1;
    cswap2['C']['S'] = 2;
    cswap2['C']['Y'] = 3;
    cswap2['G']['G'] = 0;
    cswap2['G']['R'] = 1;
    cswap2['G']['S'] = 2;
    cswap2['G']['K'] = 3;
    cswap2['T']['T'] = 0;
    cswap2['T']['W'] = 1;
    cswap2['T']['Y'] = 2;
    cswap2['T']['K'] = 3;

    haploid_models.resize(0);

    double a_prob[4] = {1.0, 0.0, 0.0, 0.0};
    double c_prob[4] = {0.0, 1.0, 0.0, 0.0};
    double g_prob[4] = {0.0, 0.0, 1.0, 0.0};
    double t_prob[4] = {0.0, 0.0, 0.0, 1.0};
    haploid_models.push_back(SnpModel( a_prob, 'A'));
    haploid_models.push_back(SnpModel( c_prob, 'C'));
    haploid_models.push_back(SnpModel( g_prob, 'G'));
    haploid_models.push_back(SnpModel( t_prob, 'T'));
    
    // The diploid models include the haploid models as a special case.
    diploid_models.resize(4);
    diploid_models[0].push_back(SnpModel( a_prob, 'A'));
    diploid_models[1].push_back(SnpModel( c_prob, 'C'));
    diploid_models[2].push_back(SnpModel( g_prob, 'G'));
    diploid_models[3].push_back(SnpModel( t_prob, 'T'));

    double m_prob[4] = {0.5, 0.5, 0.0, 0.0};
    double r_prob[4] = {0.5, 0.0, 0.5, 0.0};
    double w_prob[4] = {0.5, 0.0, 0.0, 0.5};
    double s_prob[4] = {0.0, 0.5, 0.5, 0.0};
    double y_prob[4] = {0.0, 0.5, 0.0, 0.5};
    double k_prob[4] = {0.0, 0.0, 0.5, 0.5};

    diploid_models[0].push_back(SnpModel( m_prob, 'M'));
    diploid_models[1].push_back(SnpModel( m_prob, 'M'));

    diploid_models[0].push_back(SnpModel( r_prob, 'R'));
    diploid_models[2].push_back(SnpModel( r_prob, 'R'));
    
    diploid_models[0].push_back(SnpModel( w_prob, 'W'));
    diploid_models[3].push_back(SnpModel( w_prob, 'W'));
    
    diploid_models[1].push_back(SnpModel( s_prob, 'S'));
    diploid_models[2].push_back(SnpModel( s_prob, 'S'));
    
    diploid_models[1].push_back(SnpModel( y_prob, 'Y'));
    diploid_models[3].push_back(SnpModel( y_prob, 'Y'));
    
    diploid_models[2].push_back(SnpModel( k_prob, 'K'));
    diploid_models[3].push_back(SnpModel( k_prob, 'K'));
}

// convert a SOLID color sequence with a leading base to the DNA strand it represents.
std::string color_to_seq(std::string &seq)
{
    unsigned int i;
    char b;
    std::string X(seq);
    b = seq[0];
    for(i=1; i<seq.size(); i++)
    {
        b = next_color[b][seq[i]];
        X[i] = b;
    }
    return X;
}

// Break the string str into a vector of strings, using the delimters delim as separators.
void strtok(const std::string& str, const std::string &delim, std::vector<std::string>& tokens)
{
    // Skip delimiters at beginning.
    std::string::size_type lastPos = str.find_first_not_of(delim, 0);
    // Find first "non-delimiter".
    std::string::size_type pos     = str.find_first_of(delim, lastPos);

    while (std::string::npos != pos || std::string::npos != lastPos)
    {  
        // Found a token, add it to the vector.
        tokens.push_back(str.substr(lastPos, pos - lastPos));
        // Skip delimiters.  Note the "not_of"
        lastPos = str.find_first_not_of(delim, pos);
        // Find next "non-delimiter"
        pos = str.find_first_of(delim, lastPos);
    }
}

void write_fasta(std::ostream &os, std::map<std::string, std::string> &seqs)
{
    unsigned int i;
    std::map<std::string, std::string>::iterator mssi;
    for(mssi = seqs.begin(); mssi != seqs.end(); mssi++)
    {
        os << ">" << mssi->first << std::endl;

        for(i = 0; i < mssi->second.size(); i += 70)
            os << mssi->second.substr(i, 70) << std::endl;
    }
}


void parse_fasta(const char *fname, std::map<std::string, std::string> &seqs)
{
    std::fstream f(fname);
    std::string t, seqname="", seq="";
    if(!f)
    {
        error("failed to open the fasta file: " + std::string(fname));
    }

    while(getline(f, t))
    {
        if(t[0] == '>')
        {
            if(seq != "")
            {
                if(seqs.find(seqname) != seqs.end())
                    error("while parsing the fasta file: " + std::string(fname) + "duplicate sequence names where encountered. |" + seqname + "| appears twice.");
                std::transform(seq.begin(), seq.end(),seq.begin(), ::toupper);            
                seqs[seqname] = seq;
            };
            std::stringstream ss(t.substr(1));
            ss >> seqname;
            seq = "";
        }
        else
        {
            seq += t;
        }
    }
    if(seq != "")
    {
        if(seqs.find(seqname) != seqs.end())
            error("while parsing the fasta file: " + std::string(fname) + "duplicate sequence names where encountered. |" + seqname + "| appears twice.");
        std::transform(seq.begin(), seq.end(),seq.begin(), ::toupper);            
        seqs[seqname] = seq;
    }

}

// This reads a PerM mapping file to determine the type of mapping file.
// TODO: we should come up with a better way of doing this?
// TODO: support PERM/SAM/BAM/BED.
// TODO: support zipped files.
// TODO: better method to recognize SAM. 
// TODO: use samread() to get sequence length.
mapping_type get_mapping_type(const char *fname)
{
    mapping_type mt = 0;
    std::string line;
    std::vector<std::string> toks;

    // open the file, and check for '@' or 'BAM' magic number.
    std::fstream f(fname);
    
    if(!f)
        error("failed to open mapping file to determine type: " + std::string(fname));

    char buf[4];

    f.read(buf,4);

    if(!f.good())
        error("failed to read mapping file when determining type: " + std::string(fname));

    /*
    if( strncmp(buf, "BAM\001", 4) == 0 )
    {
        mt |= BAM_FILE;
        mt |= SEQUENCE;
        return mt;
    }
    else 
    */
    if (strncmp(buf, "@", 1) == 0)
    {
        mt |= SAM_FILE;
        mt |= SEQUENCE;
        return mt;
    }
    else 
    {
        mt |= PERM_FILE;
        f.seekg(0);
        getline(f, line);

        // TODO: did getline work?
        strtok(line, "\t", toks);

        if( toks[1][1] == 'A' || toks[1][1] == 'C' || toks[1][1] == 'G' || toks[1][1] == 'T' )
            mt |= SEQUENCE;
        
        else if( toks[1][1] == '0' || toks[1][1] == '1' || toks[1][1] == '2' || toks[1][1] == '3' )
            mt |= COLOR;
        else
            error("Failed to read PerM mapping file read column.");

        return mt;
    }
}

// This reads a PerM mapping file to calculate the read length of the sequences in the mapping file.
// TODO: we should come up with a better way of doing this?
int get_read_length(const char *fname)
{
    mapping_type mt = get_mapping_type(fname);

    std::string line;
    std::vector<std::string> toks;
    std::ifstream fin(fname, std::ifstream::in);
    if(!fin)
        error("Could not open mapping file: " + std::string(fname));

    // TODO: check if gzipped?
    
    getline(fin, line);

    // TODO: did getline work?
    strtok(line, "\t", toks);

    if(mt & PERM_FILE)
    {
        return toks[1].size();
    }
    else if(mt & SAM_FILE)
    {

        samfile_t *sam_fp = samopen(fname, "r", 0);
        bam1_t *read = bam_init1();
        
        if(samread(sam_fp, read) >= 0)
        {
            int len = read->core.l_qseq;
            bam_destroy1(read);
            samclose(sam_fp);
            return len;
        }
        else
        {
            bam_destroy1(read);
            samclose(sam_fp);
            error("Problem reading SAM file: " + std::string(fname));
        }
    }

    error("Mapping Type was not PERM or SAM?");
    return 0;
}


// Constructors for the ColorRead object.
ColorRead::ColorRead(std::string _id, std::string _seq, std::vector<int> _qual, std::vector<MapLoci> _map_locs, int _num) 
    : id(_id), seq(_seq), qual(_qual), map_locs(_map_locs), num(_num) {}

ColorRead::ColorRead() {}

// Sequence Read.
SeqRead::SeqRead(std::string _id, std::string _seq, std::vector<int> _qual, std::vector<MapLoci> _map_locs, int _num) 
    : id(_id), seq(_seq), qual(_qual), map_locs(_map_locs), num(_num) {}

SeqRead::SeqRead() {}



// Constructors of a MapLoci object.
MapLoci::MapLoci(std::string _id, unsigned int _pos, std::string _seq, char _strand) 
    : id(_id), pos(_pos), seq(_seq), strand(_strand) {}


// SnpRecord class constructors.
SnpRecord::SnpRecord(unsigned int _id, std::string &_chr, unsigned int _start, unsigned int _end, std::valarray<float> &_subs, std::valarray<float> &_totl   ) 
    : id(_id), chr(_chr), start(_start), end(_end), subs(_subs), totl(_totl) {}

SnpRecord::SnpRecord(){}


// ColorSnpCall Constructors.
ColorSnpCall::ColorSnpCall(unsigned int _id, std::string &_chr, unsigned int _start, unsigned int _end, std::string _seq, std::valarray<float> &_subs, std::valarray<float> &_totl   ) 
    : id(_id), chr(_chr), start(_start), end(_end), seq(_seq), subs(_subs), totl(_totl), total_reads(0) {}

ColorSnpCall::ColorSnpCall(){}

// SeqSnpCall Constructors.
SeqSnpCall::SeqSnpCall(unsigned int _id, std::string &_chr, unsigned int _start, unsigned int _end, std::string _seq, std::valarray<float> &_subs, std::valarray<float> &_totl   ) 
    : id(_id), chr(_chr), start(_start), end(_end), seq(_seq), subs(_subs), totl(_totl), total_reads(0) {}

SeqSnpCall::SeqSnpCall(){}



// SnpModel constructors.
SnpModel::SnpModel(double p[4], char c)
{
    for(unsigned int i = 0; i < 4; i++)
        prob[i] = p[i];
    base = c;
}

SnpModel::SnpModel(){}

